from common.mongo_utils import get_db
import pandas as pd


# start_date=20220901
def get_data_from_mongo(symbol, start_date=None):
    query = {'code': symbol}
    if start_date:
        query['date'] = {'$gte': start_date}
    items = get_db()['stock_quotes'].find(query, {'factor': 1, 'open': 1, 'high': 1, 'low': 1, 'close': 1, 'volume': 1,
                                                  'date': 1,
                                                  '_id': 0})
    df = pd.DataFrame(list(items))
    return df


def get_fin_from_mongo(symbol, indicators, start_report_date=None):
    query = {'code': symbol}
    if start_report_date:
        query['end_date'] = {'$gte': start_report_date}

    filters = {'date': 1, 'period': 1, '_id': 0}
    for ind in indicators:
        filters[ind] = 1

    items = get_db()['stock_fin'].find(query, filters)
    items = list(items)

    csv = []
    for item in items:
        for ind in indicators:
            csv.append(
                {'date': item['date'], 'period': item['period'], 'field': ind, 'value': item[ind], 'symbol': symbol})

    df = pd.DataFrame(csv)

    df['date'] = df['date'].apply(
        lambda x: (pd.to_datetime(x)).strftime('%Y-%m-%d')
    )

    # 正则化
    dt = df["period"].apply(
        lambda x: (
                pd.to_datetime(x) + pd.DateOffset(days=45)
        ).date()
    )

    # 这里只是fillna
    df["date"] = df["date"].fillna(dt.astype(str))

    df["period"] = pd.to_datetime(df["period"])
    # df["period2"] = pd.to_datetime(df["period"])
    df["period"] = df["period"].apply(
        lambda x: x.year * 100 + (x.month - 1) // 3 + 1
    )
    return df


from common.scripts.dump_bin import DumpDataAll, DumpDataUpdate, DumpDataFix
from common.config import CN_CSVS, CN_DATA_TEST, PIT_CSVS

from common.scripts.dump_pit import DumpPitData


# csv的文件路径以及qlib的目录
def dump_csv_to_bin():
    dump = DumpDataFix(csv_path=CN_CSVS, qlib_dir=CN_DATA_TEST)
    dump.dump()


def dump_csv_to_pit():
    dump = DumpPitData(csv_path=PIT_CSVS, qlib_dir=CN_DATA_TEST)
    dump.dump()


if __name__ == '__main__':
    symbol = '600519.SH'

    # df = get_data_from_mongo(symbol, start_date='20200822')
    # df['factor'] = df['factor'] + 1
    # df.to_csv(str(CN_CSVS.resolve()) + '/' + symbol + '.csv', index=False)

    df = get_fin_from_mongo(symbol, indicators=['eps', 'roe'])
    print(df)
    df.to_csv(str(PIT_CSVS.resolve()) + '/' + symbol + '.csv', index=False)
    dump_csv_to_pit()

    # dump_csv_to_bin()
